package main

import (
	"fmt"
	"math"
	"net"
	"os"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

// 全局变量用于去重
var (
	seenTcUrls  = make(map[string]bool)
	seenStreams = make(map[string]bool)
)

// 获取当前活跃的网络设备索引
func getActiveDeviceIndex(devices []pcap.Interface) int {
	// 获取所有网络接口
	interfaces, err := net.Interfaces()
	if err != nil {
		return 0
	}

	// 查找有活跃IP地址且状态为UP的接口
	for _, iface := range interfaces {
		// 检查接口是否启用
		if iface.Flags&net.FlagUp == 0 {
			continue
		}
		// 跳过回环接口
		if iface.Flags&net.FlagLoopback != 0 {
			continue
		}

		// 获取接口的地址
		addrs, err := iface.Addrs()
		if err != nil || len(addrs) == 0 {
			continue
		}

		// 检查是否有有效的IP地址（非链路本地地址）
		hasValidIP := false
		for _, addr := range addrs {
			if ipnet, ok := addr.(*net.IPNet); ok {
				ip := ipnet.IP
				// 排除链路本地地址和回环地址
				if !ip.IsLoopback() && !ip.IsLinkLocalUnicast() {
					hasValidIP = true
					break
				}
			}
		}

		if !hasValidIP {
			continue
		}

		// 在pcap设备列表中查找匹配的设备
		for i, device := range devices {
			// 尝试匹配设备名称
			// Windows上设备名称可能包含GUID，需要模糊匹配
			if containsDeviceName(device.Name, iface.Name) || containsDeviceName(device.Description, iface.Name) {
				return i
			}
			// 也检查设备的IP地址是否匹配
			for _, deviceAddr := range device.Addresses {
				for _, ifaceAddr := range addrs {
					if ipnet, ok := ifaceAddr.(*net.IPNet); ok {
						if deviceAddr.IP.Equal(ipnet.IP) {
							return i
						}
					}
				}
			}
		}
	}

	return 0
}

// 辅助函数：检查设备名称是否匹配
func containsDeviceName(deviceName, ifaceName string) bool {
	if deviceName == "" || ifaceName == "" {
		return false
	}
	// 简单的包含检查
	return deviceName == ifaceName
}

func main() {
	// 获取所有网络设备
	devices, err := pcap.FindAllDevs()
	if err != nil {
		panic(err)
	}

	// 选择设备
	if len(devices) == 0 {
		panic("No network devices found")
	}

	// 获取当前活跃的网络设备作为默认值
	defaultDeviceIndex := getActiveDeviceIndex(devices)

	// 显示所有可用的网络设备
	fmt.Fprintln(os.Stderr, "Available network devices:")
	for i, device := range devices {
		prefix := "  "
		if i == defaultDeviceIndex {
			prefix = "* " // 标记默认设备
		}
		fmt.Fprintf(os.Stderr, "%s%d: %s", prefix, i, device.Name)
		if device.Description != "" {
			fmt.Fprintf(os.Stderr, " (%s)", device.Description)
		}
		if i == defaultDeviceIndex {
			fmt.Fprintf(os.Stderr, " [当前使用]")
		}
		fmt.Fprintln(os.Stderr)
	}

	// 让用户选择设备
	var deviceIndex int
	fmt.Fprintf(os.Stderr, "\nEnter device number (0-%d, default %d): ", len(devices)-1, defaultDeviceIndex)
	_, err = fmt.Fscanln(os.Stdin, &deviceIndex)
	if err != nil {
		// 如果输入错误或直接回车，使用默认的活跃设备
		deviceIndex = defaultDeviceIndex
		fmt.Fprintf(os.Stderr, "Using default device %d\n", defaultDeviceIndex)
	}

	// 验证设备索引
	if deviceIndex < 0 || deviceIndex >= len(devices) {
		fmt.Fprintf(os.Stderr, "Invalid device index, using device 0\n")
		deviceIndex = 0
	}

	device := devices[deviceIndex].Name
	fmt.Fprintf(os.Stderr, "\nCapturing on device: %s\n", device)
	if devices[deviceIndex].Description != "" {
		fmt.Fprintf(os.Stderr, "Description: %s\n", devices[deviceIndex].Description)
	}
	fmt.Fprintln(os.Stderr, "\nWaiting for RTMP traffic...")

	// 打开设备进行实时抓包
	// snaplen: 65536 (捕获完整数据包)
	// promisc: true (混杂模式)
	// timeout: pcap.BlockForever (永久等待)
	handle, err := pcap.OpenLive(device, 65536, true, pcap.BlockForever)
	if err != nil {
		panic(err)
	}
	defer handle.Close()

	// 设置BPF过滤器，只捕获RTMP端口（1935）的TCP流量
	err = handle.SetBPFFilter("tcp port 1935")
	if err != nil {
		panic(err)
	}

	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())

	packetCount := 0
	tcpPacketCount := 0
	rtmpPortPacketCount := 0
	encapsulatedCount := 0
	rtmpPacketCount := 0
	amf0Count := 0

	for packet := range packetSource.Packets() {
		packetCount++

		// 首先检查是否有TCP层
		tcpLayer := packet.Layer(layers.LayerTypeTCP)
		if tcpLayer == nil {
			continue
		}

		tcpPacketCount++
		tcp, _ := tcpLayer.(*layers.TCP)

		// 检查是否是RTMP端口（1935是标准RTMP端口）
		isRTMPPort := tcp.DstPort == 1935 || tcp.SrcPort == 1935

		if !isRTMPPort {
			continue
		}

		rtmpPortPacketCount++

		// 获取TCP payload
		payload := tcp.Payload
		if len(payload) == 0 {
			continue
		}

		// 检查payload是否以IP头开始（0x45表示IPv4）
		// 但是要排除RTMP chunk的情况
		// RTMP chunk的第一个字节：高2位是fmt，低6位是chunk stream ID
		// 如果是0x03开头，可能是RTMP握手或者chunk stream ID = 3
		firstByte := payload[0]

		// 只有当确实是IP包时才跳过（0x45 = IPv4 version 4, header length 5）
		// 并且第二个字节应该是TOS字段（通常是0x00）
		if len(payload) > 20 && firstByte == 0x45 && payload[1] == 0x00 {
			// 进一步验证：检查IP包的总长度字段
			ipTotalLength := int(payload[2])<<8 | int(payload[3])
			if ipTotalLength > 20 && ipTotalLength <= len(payload) {
				encapsulatedCount++
				continue
			}
		}

		rtmpPacketCount++

		// 尝试解析RTMP数据
		if len(payload) >= 1 {
			// 方法1: 标准解析
			found := analyzeRTMPPacket(payload)
			if found {
				amf0Count++
			}

			// 方法2: 搜索connect命令（作为备用方案）
			if !found && searchForConnectCommand(payload) {
				amf0Count++
			}

			// 方法3: 搜索FCPublish命令
			if searchForFCPublishCommand(payload) {
				amf0Count++
			}
		}
	}

}

// 搜索connect命令（备用方案）
func searchForConnectCommand(payload []byte) bool {
	// 搜索AMF0字符串 "connect" 的模式
	// 格式: 02 00 07 63 6F 6E 6E 65 63 74
	// 02 = AMF0 String type
	// 00 07 = length (7)
	// 63 6F 6E 6E 65 63 74 = "connect"
	connectPattern := []byte{0x02, 0x00, 0x07, 0x63, 0x6F, 0x6E, 0x6E, 0x65, 0x63, 0x74}

	for i := 0; i <= len(payload)-len(connectPattern); i++ {
		// 检查是否匹配connect模式
		match := true
		for j := 0; j < len(connectPattern); j++ {
			if payload[i+j] != connectPattern[j] {
				match = false
				break
			}
		}

		if match {
			// 尝试从这个位置解析AMF0命令
			return parseAMF0Command(payload[i:])
		}
	}

	return false
}

// 搜索FCPublish命令
func searchForFCPublishCommand(payload []byte) bool {
	// 搜索AMF0字符串 "FCPublish" 的模式
	// 格式: 02 00 09 46 43 50 75 62 6C 69 73 68
	// 02 = AMF0 String type
	// 00 09 = length (9)
	// 46 43 50 75 62 6C 69 73 68 = "FCPublish"
	fcPublishPattern := []byte{0x02, 0x00, 0x09, 0x46, 0x43, 0x50, 0x75, 0x62, 0x6C, 0x69, 0x73, 0x68}

	for i := 0; i <= len(payload)-len(fcPublishPattern); i++ {
		// 检查是否匹配FCPublish模式
		match := true
		for j := 0; j < len(fcPublishPattern); j++ {
			if payload[i+j] != fcPublishPattern[j] {
				match = false
				break
			}
		}

		if match {
			// 尝试从这个位置解析AMF0命令
			return parseAMF0FCPublishCommand(payload[i:])
		}
	}

	return false
}

// 分析RTMP数据包
func analyzeRTMPPacket(payload []byte) bool {
	if len(payload) == 0 {
		return false
	}

	firstByte := payload[0]

	// 提取fmt (高2位)
	chunkFmt := (firstByte >> 6) & 0x03

	// 根据fmt确定header长度
	var headerSize int
	switch chunkFmt {
	case 0:
		headerSize = 11
	case 1:
		headerSize = 7
	case 2:
		headerSize = 3
	case 3:
		headerSize = 0
	}

	// 尝试查找RTMP消息类型
	if headerSize > 0 && len(payload) > headerSize {
		// 如果是完整的Type 0 header，解析更多信息
		if chunkFmt == 0 && len(payload) >= 12 {
			messageTypeID := payload[7]

			// 如果是AMF0命令消息（类型20），尝试解析
			if messageTypeID == 20 && len(payload) > 12 {
				amfData := payload[12:]
				return parseAMF0Command(amfData)
			}
		}
	}
	return false
}

// 解析AMF0命令消息（connect命令）
func parseAMF0Command(data []byte) bool {
	if len(data) == 0 {
		return false
	}

	offset := 0

	// 第一个字段通常是命令名称（String类型）
	if offset < len(data) && data[offset] == 0x02 {
		_, newOffset := parseAMF0String(data, offset)
		offset = newOffset

		// 第二个字段通常是Transaction ID（Number类型）
		if offset < len(data) && data[offset] == 0x00 {
			_, newOffset := parseAMF0Number(data, offset)
			offset = newOffset
		}

		// 第三个字段通常是Command Object（Object类型）
		if offset < len(data) && data[offset] == 0x03 {
			properties, _ := parseAMF0Object(data, offset)

			// 只输出tcUrl的值（去重）
			if tcUrl, exists := properties["tcUrl"]; exists {
				if !seenTcUrls[tcUrl] {
					seenTcUrls[tcUrl] = true
					fmt.Println("服务器：" + tcUrl)
				}
				return true
			}
		}
	}
	return false
}

// 解析AMF0 FCPublish命令
func parseAMF0FCPublishCommand(data []byte) bool {
	if len(data) == 0 {
		return false
	}

	offset := 0

	// 第一个字段是命令名称 "FCPublish"（String类型）
	if offset < len(data) && data[offset] == 0x02 {
		_, newOffset := parseAMF0String(data, offset)
		offset = newOffset

		// 第二个字段是Transaction ID（Number类型）
		if offset < len(data) && data[offset] == 0x00 {
			_, newOffset := parseAMF0Number(data, offset)
			offset = newOffset
		}

		// 第三个字段可能是Null（0x05）
		if offset < len(data) && data[offset] == 0x05 {
			offset++
		}

		// 第四个字段是stream name（String类型）
		if offset < len(data) && data[offset] == 0x02 {
			streamName, _ := parseAMF0String(data, offset)
			if streamName != "" {
				// 去重输出
				if !seenStreams[streamName] {
					seenStreams[streamName] = true
					fmt.Println("推流码：" + streamName)
				}
				return true
			}
		}
	}
	return false
}

// 解析AMF0字符串
func parseAMF0String(data []byte, offset int) (string, int) {
	if offset+3 > len(data) {
		return "", offset
	}

	// 跳过类型标记（0x02）
	offset++

	// 读取字符串长度（2字节，大端序）
	length := int(data[offset])<<8 | int(data[offset+1])
	offset += 2

	if offset+length > len(data) {
		return "", offset
	}

	// 读取字符串内容
	str := string(data[offset : offset+length])
	offset += length

	return str, offset
}

// 解析AMF0数字
func parseAMF0Number(data []byte, offset int) (float64, int) {
	if offset+9 > len(data) {
		return 0, offset
	}

	// 跳过类型标记（0x00）
	offset++

	// 读取8字节双精度浮点数（大端序）
	bits := uint64(data[offset])<<56 | uint64(data[offset+1])<<48 |
		uint64(data[offset+2])<<40 | uint64(data[offset+3])<<32 |
		uint64(data[offset+4])<<24 | uint64(data[offset+5])<<16 |
		uint64(data[offset+6])<<8 | uint64(data[offset+7])

	offset += 8

	// 将位模式转换为float64
	number := math.Float64frombits(bits)

	return number, offset
}

// 解析AMF0对象
func parseAMF0Object(data []byte, offset int) (map[string]string, int) {
	properties := make(map[string]string)

	if offset >= len(data) {
		return properties, offset
	}

	// 跳过类型标记（0x03）
	offset++

	// 循环读取属性，直到遇到对象结束标记（0x00 0x00 0x09）
	for offset < len(data) {
		// 检查对象结束标记
		if offset+3 <= len(data) && data[offset] == 0x00 && data[offset+1] == 0x00 && data[offset+2] == 0x09 {
			offset += 3
			break
		}

		// 读取属性名（UTF-8字符串，2字节长度 + 内容）
		if offset+2 > len(data) {
			break
		}

		nameLength := int(data[offset])<<8 | int(data[offset+1])
		offset += 2

		if offset+nameLength > len(data) {
			break
		}

		propertyName := string(data[offset : offset+nameLength])
		offset += nameLength

		// 读取属性值（根据类型）
		if offset >= len(data) {
			break
		}

		valueType := data[offset]
		var propertyValue string

		switch valueType {
		case 0x00: // Number
			number, newOffset := parseAMF0Number(data, offset)
			propertyValue = fmt.Sprintf("%.0f", number)
			offset = newOffset

		case 0x01: // Boolean
			if offset+2 <= len(data) {
				offset++
				if data[offset] != 0 {
					propertyValue = "true"
				} else {
					propertyValue = "false"
				}
				offset++
			}

		case 0x02: // String
			str, newOffset := parseAMF0String(data, offset)
			propertyValue = str
			offset = newOffset

		case 0x05: // Null
			propertyValue = "null"
			offset++

		case 0x06: // Undefined
			propertyValue = "undefined"
			offset++

		default:
			// 未知类型，跳过
			offset++
		}

		properties[propertyName] = propertyValue
	}

	return properties, offset
}
